// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once
#include <algorithm>
#include <iostream>
#include <map>
#include <memory>
#include <random>

#include "coin2/exchange/base/feed/update.h"
#include "coin2/exchange/base/symbology/product_encyclopedia.h"
#include "coin2/strategy/strategy.h"
#include "presto/quant/math/moving_average.h"
#include "presto/quant/math/moving_window.h"

using coin::proto::OrderDirection;
using coin::proto::OrderDuration;
using coin::proto::OrderType;
using coin2::exchange::base::feed::FeedUpdate;
using ::coin2::exchange::base::symbology::ProductEncyclopedia;
using presto::math::MovingWindowWithCallee;
using presto::math::TimeWindowMovingAverage;
using presto::math::TimeWindowMovingMaximum;
using namespace coin2::exchange::base::symbology;
using namespace coin2::base::config;
using namespace coin2::exchange::feed;
using namespace coin2::app;

enum Features {
  // trade
  Vwap5s,
  Vwap15s,
  Vwap30s,
  Vwas5s,  // volume-weighted average sign
  Vwas15s,
  Vwas30s,
  TradeQty5s,  // sum(trade.unsigned_qty)
  TradeQty15s,
  TradeQty30s,
  // book
  Midp,
  Midp5s,
  Midp15s,
  Midp30s,
  Ask0p,
  Ask0q,
  Ask1p,
  Ask1q,
  Ask2p,
  Ask2q,
  Bid0p,
  Bid0q,
  Bid1p,
  Bid1q,
  Bid2p,
  Bid2q,
  SpreadMin5s,
  SpreadMax5s,
  //
  NUM_FEATURE,
};

enum Y {
  Ready,
  AskToBid3s,
  AskToBid5s,
  AskToBid10s,
  AskToBid30s,
  BidToAsk3s,
  BidToAsk5s,
  BidToAsk10s,
  BidToAsk30s,
  //
  NUM_Y,
};

const int64_t TRADE_VALID_PERIOD = 20'000'000;       // 20ms
const int64_t BOOK_VALID_PERIOD = 2'000'000'000;     // 2s
const int64_t SELF_BOOK_VALID_PERIOD = 500'000'000;  // 500ms
// const int64_t COOLTIME = 500'000'000;                // 500ms
// const int64_t MIN = 60 * 1'000'000'000LL;
const int64_t SEC = 1'000'000'000LL;

class ProductModel : public ::coin2::strategy::IStrategy {
 public:
  ProductModel(const IProduct& product, std::array<double, Features::NUM_FEATURE>* features_out)
      : product_(product), features_(*features_out), spread_5s_(0) {
    (void)product_;

    midp_5s_.ResizeWindow(5 * SEC);
    midp_15s_.ResizeWindow(15 * SEC);
    midp_30s_.ResizeWindow(30 * SEC);

    trade_qty_5s_.ResizeWindow(5 * SEC);
    trade_qty_15s_.ResizeWindow(15 * SEC);
    trade_qty_30s_.ResizeWindow(30 * SEC);

    pq_5s_.ResizeWindow(5 * SEC);
    pq_15s_.ResizeWindow(15 * SEC);
    pq_30s_.ResizeWindow(30 * SEC);

    spread_5s_.ResizeWindow(5 * SEC);

    for (int i = 0; i < Features::NUM_FEATURE; i++) {
      features_[i] = 0;
    }
  }

  bool IsReady(int64_t now) const {
    return (now - last_trade_ < TRADE_VALID_PERIOD) && (now - last_book_ < BOOK_VALID_PERIOD);
  }

  bool IsReadyAsTarget(int64_t now) const {
    return (now - last_trade_ < TRADE_VALID_PERIOD) && (now - last_book_ < SELF_BOOK_VALID_PERIOD);
  }

  void onTradeFeed(const FeedUpdate& upd) override {
    auto ts = upd.timestamp();
    last_trade_ = ts;
    // update vwap
    // turnover
    const auto& trade = upd.trade();
    trade_qty_5s_.Update(ts, trade.fill_qty);
    trade_qty_15s_.Update(ts, trade.fill_qty);
    trade_qty_30s_.Update(ts, trade.fill_qty);

    double signed_qty = trade.signed_qty();
    trade_qty_signed_5s_.Update(ts, signed_qty);
    trade_qty_signed_15s_.Update(ts, signed_qty);
    trade_qty_signed_30s_.Update(ts, signed_qty);

    pq_5s_.Update(ts, trade.price * trade.fill_qty);
    pq_15s_.Update(ts, trade.price * trade.fill_qty);
    pq_30s_.Update(ts, trade.price * trade.fill_qty);

    features_[Features::Vwap5s] =
        (trade_qty_5s_.sum() > 0) ? pq_5s_.sum() / trade_qty_5s_.sum() : 0;
    features_[Features::Vwap15s] =
        (trade_qty_15s_.sum() > 0) ? pq_15s_.sum() / trade_qty_15s_.sum() : 0;
    features_[Features::Vwap30s] =
        (trade_qty_30s_.sum() > 0) ? pq_30s_.sum() / trade_qty_30s_.sum() : 0;

    features_[Features::Vwas5s] = trade_qty_signed_5s_.average();
    features_[Features::Vwas15s] = trade_qty_signed_15s_.average();
    features_[Features::Vwas30s] = trade_qty_signed_30s_.average();

    features_[Features::TradeQty5s] = trade_qty_5s_.sum();
    features_[Features::TradeQty15s] = trade_qty_15s_.sum();
    features_[Features::TradeQty30s] = trade_qty_30s_.sum();
  }

  void onBookFeed(const FeedUpdate& upd) override {
    auto ts = upd.timestamp();
    const auto& book = upd.book();

    // L2
    if (book.AskN(2) && book.BidN(2)) {
      last_book_ = upd.timestamp();
      midp_ = (book.Ask0()->price + book.Bid0()->price) / 2;

      midp_5s_.Update(ts, midp_);
      midp_15s_.Update(ts, midp_);
      midp_30s_.Update(ts, midp_);
      spread_5s_.Update(ts, book.Ask0()->price - book.Bid0()->price);

      features_[Features::Midp] = midp_;
      features_[Features::Midp5s] = midp_5s_.average();
      features_[Features::Midp15s] = midp_15s_.average();
      features_[Features::Midp30s] = midp_30s_.average();
      features_[Features::Ask0p] = book.Ask0()->price;
      features_[Features::Ask0q] = book.Ask0()->qty;
      features_[Features::Ask1p] = book.AskN(1)->price;
      features_[Features::Ask1q] = book.AskN(1)->qty;
      features_[Features::Ask2p] = book.AskN(2)->price;
      features_[Features::Ask2q] = book.AskN(2)->qty;
      features_[Features::Bid0p] = book.Bid0()->price;
      features_[Features::Bid0q] = book.Bid0()->qty;
      features_[Features::Bid1p] = book.BidN(1)->price;
      features_[Features::Bid1q] = book.BidN(1)->qty;
      features_[Features::Bid2p] = book.BidN(2)->price;
      features_[Features::Bid2q] = book.BidN(2)->qty;
      spread_5s_.GetMinElement(&features_[Features::SpreadMin5s]);
      spread_5s_.GetMaxElement(&features_[Features::SpreadMax5s]);
    }
  }

  const IProduct& product_;
  std::array<double, Features::NUM_FEATURE>& features_;

  int64_t last_trade_ = 0;
  int64_t last_book_ = 0;

  // trade
  TimeWindowMovingAverage<double> trade_qty_5s_;
  TimeWindowMovingAverage<double> trade_qty_15s_;
  TimeWindowMovingAverage<double> trade_qty_30s_;
  TimeWindowMovingAverage<double> trade_qty_signed_5s_;
  TimeWindowMovingAverage<double> trade_qty_signed_15s_;
  TimeWindowMovingAverage<double> trade_qty_signed_30s_;
  TimeWindowMovingAverage<double> pq_5s_;
  TimeWindowMovingAverage<double> pq_15s_;
  TimeWindowMovingAverage<double> pq_30s_;

  // book
  double midp_ = 0;
  TimeWindowMovingAverage<double> midp_5s_;
  TimeWindowMovingAverage<double> midp_15s_;
  TimeWindowMovingAverage<double> midp_30s_;
  TimeWindowMovingMaximum<double> spread_5s_;
};

class MultiProductModel : public ::coin2::strategy::IStrategy {
 public:
  void Init(::coin2::app::Driver* driver) override {
    // feed products
    auto feed_products = driver->feed()->GetProducts();
    feature_size_ = Features::NUM_FEATURE * feed_products.size();
    features_.reset(new double[feature_size_]);

    // only one product for order
    auto order_products = driver->order()->GetProducts();
    CHECK_EQ_THROW(order_products.size(), 1);
    target_ = const_cast<coin2::exchange::base::symbology::IProduct *>(order_products[0]);

    {
      bool found = false;
      for (auto* product : feed_products) {
        if (*product == *target_) {
          found = true;
          break;
        }
      }
      if (!found) {
        throw std::runtime_error(
            "target product " + target_->absolute_norm() +
            " 's feed is not subscribed (num_feed_produces=" +
            std::to_string(feed_products.size()) + ")");
      }
    }

    // create feature model for each product
    // e.g.
    // features_[ 0~25] for product1
    // features_[25~50] for product2
    int offset = 0;
    for (auto* p : feed_products) {
      product_models_[p->key()] = std::make_unique<ProductModel>(
          *p,
          reinterpret_cast<std::array<double, Features::NUM_FEATURE>*>(&features_.get()[offset]));
      if (p->key() == target_->key()) {
        target_model_offset_ = offset;
      }
      offset += Features::NUM_FEATURE;
    }

    target_model_ = product_models_[target_->key()].get();
  }

  bool IsReady(int64_t now) {
    for (auto& [key, model] : product_models_) {
      (void)key;
      if (!model->IsReady(now)) {
        return false;
      }
    }
    return target_model_->IsReadyAsTarget(now);
  }

  void onTradeFeed(const FeedUpdate& upd) override {
    auto key = upd.product().key();
    for (auto& [key2, model] : product_models_) {
      if (key == key2) {
        model->onTradeFeed(upd);
      }
    }
  }
  void onBookFeed(const FeedUpdate& upd) override {
    auto key = upd.product().key();
    for (auto& [key2, model] : product_models_) {
      if (key == key2) {
        model->onBookFeed(upd);
      }
    }
  }

  size_t feature_size_ = 0;
  std::unique_ptr<double[]> features_;
  std::map<IProduct::TKey, std::unique_ptr<ProductModel>> product_models_;
  IProduct* target_ = nullptr;
  ProductModel* target_model_ = nullptr;
  std::ptrdiff_t target_model_offset_ = 0;
  ProductEncyclopedia product_cache_;
};

std::random_device rd;
std::mt19937_64 mt(rd());
std::uniform_int_distribution<uint64_t> dist;

struct ThrottleSampler {
  int64_t t = 0;
  const int64_t mn = 1;
  const int64_t mx = 10 * 60 * 1000000000LL;
  bool sample(int64_t ts) {
    if (t < ts) {
      int64_t delta = dist(mt) % (mx - mn + 1) + mn;
      t = ts + delta;
      return true;
    }
    return false;
  }
};

struct TradeQtySampler : ThrottleSampler {
  // XRP 500
  //  3982003 / 50000 = 79
  // 12012425 / 50000 = 240
  static constexpr double target_trade_qty_5s = 500.0;
  bool sample(int64_t ts, double trade_qty_5s) {
    double p = std::min(1.0, trade_qty_5s / target_trade_qty_5s) * 0.9;
    double r = static_cast<double>(rand()) / RAND_MAX;
    // higher trade_qty, higher prob
    if (r <= p && ThrottleSampler::sample(ts)) {
      return true;
    }
    return false;
  }
};